import csv
import os

import mysql.connector
from mysql.connector import Error
from shapely import wkb
import glob

def import_csv_to_mysql(csv_file_path, host, database, user, password):
    try:
        # 读取CSV文件
        with open(csv_file_path, mode='r', encoding='gbk') as file:
            csv_reader = csv.reader(file)
            headers = next(csv_reader)  # 获取CSV文件的列标题
            print(f"\n{csv_file_path} CSV文件列标题:")
            for i, header in enumerate(headers):
                print(f"{i}: {header}")

            geom_index = headers.index('geom')
            address_index = headers.index('address')
            address_name_index = headers.index('address_name')
            # 询问用户映射关系
            # address_col = int(input(f"\n请输入要导入到address字段的列号: "))
            # address_name_col = int(input(f"\n请输入要导入到address_name字段的列号: "))
            # location_col = int(input("请输入要导入到location字段的列号: "))
            address_col = address_index
            address_name_col = address_name_index
            location_col = geom_index
            # 连接到MySQL数据库
            connection = mysql.connector.connect(
                host=host,
                database=database,
                user=user,
                password=password
            )

            if connection.is_connected():
                cursor = connection.cursor()

                # 准备插入语句
                insert_query = """INSERT INTO gis_poi_data (address, address_name, location)
                                  VALUES (%s, %s, %s)"""

                # 读取并插入数据
                for row in csv_reader:
                    try:
                        address = row[address_col] if address_col < len(row) else None
                        address_name = row[address_name_col] if address_name_col < len(row) else None
                        location = row[location_col] if location_col < len(row) else None
                        if "|" in location or "||" in location:
                            location = location.split("|")[0]
                        if "point" not in location.lower():
                            geometry = wkb.loads(location)
                            location = geometry.wkt

                        # 执行插入
                        cursor.execute(insert_query, (address.strip(), address_name.strip(), location.strip()))
                    except IndexError:
                        print(f"跳过行 {csv_reader.line_num}: 列索引超出范围")
                        continue

                # 提交事务
                connection.commit()
                print(f"\n成功导入 {csv_reader.line_num - 1} 行数据到MySQL数据库")

    except Error as e:
        print(f"数据库错误: {e}")
    except Exception as e:
        print(f"发生错误: {e}")
    finally:
        if connection and connection.is_connected():
            cursor.close()
            connection.close()
            print("MySQL连接已关闭")


# 使用示例
if __name__ == "__main__":
    # 输入参数
    # csv_path = input("请输入CSV文件路径: ")
    csv_path = r"D:\测试数据\重要场所poi"
    # host = input("MySQL主机地址: 127.0.0.1")
    # database = input("数据库名称: xz")
    # user = input("用户名: root")
    # password = input("密码: egova")

    for path in glob.glob(os.path.join(csv_path, "*.csv")):
        # 调用函数
        import_csv_to_mysql(path, "127.0.0.1", "xz", "root", "egova")
